import os
import numpy as np
import re


def extract_prefix(filename):
    pattern = re.compile(r'(.*)tr\d+-tc\d+\.npy')
    match = pattern.match(filename)
    if match:
        return match.group(1)
    return ""


def cut_and_save_large_flow(flow_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    flow_files = sorted([f for f in os.listdir(flow_dir) if f.endswith('.npy')])

    for flow_file in flow_files:
        prefix = extract_prefix(flow_file)
        # 解析行和列编号
        match = re.search(r'tr(\d+)-tc(\d+)', flow_file)
        if not match:
            continue

        big_tr = int(match.group(1))
        big_tc = int(match.group(2))

        flow_path = os.path.join(flow_dir, flow_file)
        large_flow = np.load(flow_path)
        if large_flow is None:
            raise FileNotFoundError(f"Unable to open flow file: {flow_path}")

        large_tile_size = 1024

        # 计算大光流场的起始和结束小光流场编号
        if big_tr == 1:
            row_start, row_end = 1, 15
            row_valid_start, row_valid_end = 1, 13
        elif big_tr == 2:
            row_start, row_end = 12, 26
            row_valid_start, row_valid_end = 14, 24
        else:
            row_start, row_end = 23, 37
            row_valid_start, row_valid_end = 25, 37

        if big_tc == 1:
            col_start, col_end = 1, 13
            col_valid_start, col_valid_end = 1, 12
        elif big_tc == 2:
            col_start, col_end = 11, 23
            col_valid_start, col_valid_end = 13, 21
        else:
            col_start, col_end = 21, 33
            col_valid_start, col_valid_end = 22, 33

        for small_row in range(row_start, row_end + 1):
            if small_row < row_valid_start or small_row > row_valid_end:
                continue
            for small_col in range(col_start, col_end + 1):
                if small_col < col_valid_start or small_col > col_valid_end:
                    continue

                x_offset = (small_row - row_start) * large_tile_size
                y_offset = (small_col - col_start) * large_tile_size

                small_flow = large_flow[x_offset:x_offset + large_tile_size, y_offset:y_offset + large_tile_size]

                small_filename = f"{prefix}tr{small_row}-tc{small_col}.npy"
                small_flow_path = os.path.join(output_dir, small_filename)
                np.save(small_flow_path, small_flow)
                print(f"Saved {small_filename} to {small_flow_path}")

# 示例用法
flow_dir = '/root/data/deep_align/32nm_flow_big(3)'  # 替换为实际的大光流场文件夹路径
output_dir = '/root/data/deep_align/32nm_flow(3)'  # 替换为实际的小光流场输出文件夹路径

cut_and_save_large_flow(flow_dir, output_dir)
